from src.E import ModelG
from src.D import ModelD
from src.TOOL import weight_init
import torch.optim as optim
from torch.autograd import Variable
import torch
import torch.nn as nn
from src.datasetting import setting
import numpy as np
from src.TOOL import make_show
import torchvision

batch_size=64
input_dim=100
num_epochs=10

#训练准备
netG=ModelG()
netG.apply(weight_init)
netG=netG.cuda()

netD=ModelD()
netD.apply(weight_init)
netD=netD.cuda()

dtype=torch.cuda.FloatTensor
itype=torch.cuda.LongTensor

optimizerD=optim.Adam(netD.parameters(),lr=0.0002,betas=(0.5,0.999))
optimizerG=optim.Adam(netG.parameters(),lr=0.0002,betas=(0.5,0.999))

noise=Variable(dtype(batch_size,input_dim,1,1))
fixed_noise=Variable(dtype(batch_size,input_dim,1,1).normal_(0,1))
noise=noise.cuda()
fixed_noise=fixed_noise.cuda()


criterion=nn.BCELoss()
error_G=None #总误差

results=[]

train_loader,val_loader,test_loader=setting()

for epoch in range(num_epochs):
    for batch_idx, (data, _) in enumerate(train_loader):

        # 判别器
        ##真图判别
        optimizerD.zero_grad()
        data = Variable(data)
        data = data.cuda()

        labellen=data.size()[0]
        label = Variable(torch.ones(labellen))#data.size()[0]为batch_size
        label = label.cuda()



        netD.train()
        output = netD(data)
        output = torch.squeeze(output)

        # label.data.fill_(1)

        error_real = criterion(output, label)#和真图对比
        error_real.backward()
        # D_x=output.data.mean()

        ##假图判别
        with torch.no_grad():
            noise.resize_(data.size()[0], input_dim, 1, 1)
        noise.normal_(0, 1)
        fake_pic = netG(noise).detach()  # detach是为了不让生成器参与梯度更新
        output2 = netD(fake_pic)
        label.data.fill_(0)  # 标志假图
        output2 = torch.squeeze(output2)
        error_fake = criterion(output2, label)
        error_fake.backward()
        error_D = error_fake + error_real
        error_D = error_D.cpu()
        optimizerD.step()

        # 生成器
        if error_G is None or np.random.rand() < 0.5:
            optimizerG.zero_grad()

            # 生成器的目标函数与判别器的是相反的，因此当判别器无法辨别的时候为正确
            label.data.fill_(1)
            # 随机产生噪声
            noise.data.normal_(0, 1)
            netG.train()
            fake_pic = netG(noise)
            output = netD(fake_pic)
            output=torch.squeeze(output)
            error_G = criterion(output, label)
            error_G.backward()
            error_G = error_G.cpu()
            optimizerG.step()

        results.append([float(error_D.data.numpy()), float(error_G.data.numpy())])

        if batch_idx % 100 == 0:
            print('第{}周期,进度：{}/{},判别器损失:{:.2f},生成器损失:{:.2f}'.format(epoch, batch_idx, len(train_loader),
                                                                    error_D.item(), error_G.item()))

    # 一轮生成64个图片 输出到文件
    netG.eval()
    fake_u = netG(fixed_noise)
    img = make_show(fake_u)
    torchvision.utils.save_image(img, '../img/第%s周期假图.jpg' % (epoch))

    data, _ = next(iter(train_loader))
    img = make_show(Variable(data))
    torchvision.utils.save_image(img, '../img/第%s周期真图.jpg' % (epoch))